'''
This is a difficult activity which will involve some relatively challenging filtering/sorting of data
You will want to have the docs or a guide on using matplotlib/pandas available, it is not expected that you do this completely independently.
chatGPT is very good when debugging Pandas, as it decodes error messages into more readable explanations- so if you are struggling to understand what 
you have done wrong I recommend pasting error messages into it!

All the basic program code has been included, you should be able to simply write the code to filter the data into an appropriate dataframe and then
plot and show the graphs.
'''

import pandas as pd
import matplotlib.pyplot as plt


def average_bar_chart(data):
    '''
    Complete this code to produce a bar chart of the total of each 'hours per day' column
    '''
    # Filter for columns that start with 'Hours_'
    hours_columns = [col for col in data.columns if col.startswith('Hours_')]
    
    # Calculate the average for each hours column
    averages = data[hours_columns].mean()
    
    # Create the bar chart
    plt.figure(figsize=(10, 6))
    averages.plot(kind='bar', color='skyblue', edgecolor='black')
    plt.title('Average Hours Per Day for Different Activities')
    plt.xlabel('Activity')
    plt.ylabel('Average Hours')
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.show()

def hours_histogram(data, column_name):
    '''
    Complete this code so that it can display a histogram of any given 'hours per day' column
    '''
    # Check if the column exists in the data
    if column_name not in data.columns:
        print(f"Column '{column_name}' not found in the data.")
        return
    
    # Create histogram
    plt.figure(figsize=(8, 6))
    plt.hist(data[column_name].dropna(), bins=10, color='lightgreen', edgecolor='black', alpha=0.7)
    plt.title(f'Distribution of {column_name}')
    plt.xlabel('Hours')
    plt.ylabel('Frequency')
    plt.tight_layout()
    plt.show()

def hours_pie_chart(data, student_id):
    '''
    Complete this code so that it produces a pie chart of any given student's hours per day - this should be selected by student ID (the first column)
    '''
    # Find the student by ID
    student_data = data[data['Student_ID'] == student_id]
    
    # Check if student exists
    if student_data.empty:
        print(f"Student with ID {student_id} not found.")
        return
    
    # Get hours columns
    hours_columns = [col for col in data.columns if col.startswith('Hours_')]
    
    # Extract the hours values for this student
    student_hours = student_data[hours_columns].iloc[0]
    
    # Create pie chart
    plt.figure(figsize=(8, 8))
    plt.pie(student_hours, labels=[col.replace('Hours_', '') for col in hours_columns], 
            autopct='%1.1f%%', startangle=90)
    plt.title(f'Time Distribution for Student {student_id}')
    plt.tight_layout()
    plt.show()

def total_hours_per_stress_level(data, stress_level):
    '''
    Complete this code to create a stacked bar chart showing a breakdown of the total hours worked by students, grouped by their reported stress level.
    '''
    # Filter data by stress level
    stress_data = data[data['Stress_Level'] == stress_level]
    
    # Check if any data exists for this stress level
    if stress_data.empty:
        print(f"No data found for stress level '{stress_level}'")
        return
    
    # Get hours columns
    hours_columns = [col for col in data.columns if col.startswith('Hours_')]
    
    # Calculate total hours for each activity
    total_hours = stress_data[hours_columns].sum()
    
    # Create stacked bar chart
    plt.figure(figsize=(10, 6))
    total_hours.plot(kind='bar', color=['red', 'blue', 'green', 'orange', 'purple', 'brown'])
    plt.title(f'Total Hours by Activity for Stress Level: {stress_level}')
    plt.xlabel('Activity')
    plt.ylabel('Total Hours')
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.show()

def display_menu():
    print("Student Lifestyle Data Analysis")
    print("1 - view average hours per day")
    print("2 - view specific activity")
    print("3 - view per-student breakdown")
    print("4 - view per stress level breakdown")
    print("Q - quit")
    choice = 0
    while choice not in ["1","2","3","4","Q"]:
        choice = input("Select option: ").upper()
    return choice

## main program:

data = pd.read_csv("lifestyle_data.csv")

# Note - when you run this in Interactive mode, the 'input' will be through a pop up instead of the terminal.
while True:
    choice = display_menu()
    match choice:
        case "Q":
            break
        case "1":
            average_bar_chart(data)
        case "2":
            # Show available hours columns to help user
            hours_columns = [col for col in data.columns if col.startswith('Hours_')]
            print(f"Available columns: {', '.join(hours_columns)}")
            col = input("Enter Hours Per column to analyse: ")
            hours_histogram(data, col)
        case "3":
            # Show available student IDs to help user
            print(f"Student IDs range from {data['Student_ID'].min()} to {data['Student_ID'].max()}")
            stid = int(input("Enter student ID: "))
            hours_pie_chart(data, stid)
        case "4":
            # Show available stress levels to help user
            stress_levels = data['Stress_Level'].unique()
            print(f"Available stress levels: {', '.join(stress_levels)}")
            stress_level = input("Enter stress level: ").title()
            total_hours_per_stress_level(data, stress_level)